/*
 * Copyright 2021 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.room.compiler.processing

import javax.tools.Diagnostic

/**
 * Common interface for basic annotation processors.
 *
 * A processor should not implement this interface directly and instead should extend
 * [androidx.room.compiler.processing.javac.JavacBasicAnnotationProcessor] or
 * [androidx.room.compiler.processing.ksp.KspBasicAnnotationProcessor].
 *
 * The XProcessing Javac and KSP implementations of this interface will automatically validate and
 * defer annotated elements for the steps, unless disabled via
 * [XProcessingEnvConfig.disableAnnotatedElementValidation]. If validation is enable and no valid
 * annotated element is found for a [XProcessingStep] then its [XProcessingStep.process] function
 * will not be invoked, except for the last round (regardless if the annotated elements are valid or
 * not). If there were invalid annotated elements until the last round, then the XProcessing
 * implementations will report an error for each invalid element. If validation is disabled, no
 * error is reported if there are invalid elements found in the ast round.
 *
 * Be aware that even though the similarity in name, the Javac implementation of this interface is
 * not 1:1 with [com.google.auto.common.BasicAnnotationProcessor]. Specifically, when validation is
 * enabled it is done for each annotated element as opposed to the enclosing type element of the
 * annotated elements for the [XProcessingStep].
 */
interface XBasicAnnotationProcessor {

    /** Returns the [XProcessingEnv]. */
    val xProcessingEnv: XProcessingEnv

    /**
     * Initializes the processor with the processing environment.
     *
     * This will be invoked before any other function in this interface and before any
     * [processingSteps].
     */
    fun initialize(env: XProcessingEnv) {}

    /** Called at the beginning of a processing round before all [processingSteps] execute. */
    fun preRound(env: XProcessingEnv, round: XRoundEnv) {}

    /** The list of processing steps to execute. */
    fun processingSteps(): Iterable<XProcessingStep>

    /** Called at the end of a processing round after all [processingSteps] have been executed. */
    fun postRound(env: XProcessingEnv, round: XRoundEnv) {}
}

/** Common code for implementations of [XBasicAnnotationProcessor] offered by XProcessing. */
internal class CommonProcessorDelegate(
    private val processorClass: Class<*>,
    private val env: XProcessingEnv,
    private val steps: List<XProcessingStep>,
) {
    // Type names of deferred elements from the processor.
    private val deferredElementNames = mutableSetOf<String>()
    // Type element names containing deferred elements from processing steps.
    private val elementsDeferredBySteps = mutableMapOf<XProcessingStep, Set<String>>()

    fun processRound(roundEnv: XRoundEnv) {
        val previousRoundDeferredElementNames = deferredElementNames.toMutableSet()
        deferredElementNames.clear()
        val currentElementsDeferredByStep =
            steps.associateWith { step ->
                // Previous round processor deferred elements, these need to be re-validated.
                val previousRoundDeferredElementsByAnnotation =
                    getStepElementsByAnnotation(step, previousRoundDeferredElementNames)
                        .withDefault { emptySet() }
                // Previous round step deferred elements, these don't need to be re-validated.
                val stepDeferredElementsByAnnotation =
                    getStepElementsByAnnotation(step, elementsDeferredBySteps[step] ?: emptySet())
                        .withDefault { emptySet() }
                val deferredElements = mutableSetOf<XElement>()
                val elementsByAnnotation =
                    step
                        .annotations()
                        .mapNotNull { annotation ->
                            val annotatedElements =
                                roundEnv.getElementsAnnotatedWith(annotation) +
                                    previousRoundDeferredElementsByAnnotation.getValue(annotation)
                            // Split between valid and invalid elements. Unlike auto-common,
                            // validation is only
                            // done in the annotated element from the round and not in the closest
                            // enclosing
                            // type element.
                            val (validElements, invalidElements) =
                                if (env.config.disableAnnotatedElementValidation) {
                                    annotatedElements to emptySet<XElement>()
                                } else {
                                    annotatedElements.partition {
                                        it.closestMemberContainer.validate()
                                    }
                                }
                            deferredElements.addAll(invalidElements)
                            (validElements + stepDeferredElementsByAnnotation.getValue(annotation))
                                .let {
                                    if (it.isNotEmpty()) {
                                        annotation to it.toSet()
                                    } else {
                                        null
                                    }
                                }
                        }
                        .toMap()
                // Store all processor deferred elements.
                deferredElementNames.addAll(
                    deferredElements.mapNotNull {
                        (it.closestMemberContainer as? XTypeElement)?.qualifiedName
                    }
                )
                // Only process the step if there are annotated elements found for this step, or the
                // step supports processing all annotations "*".
                val supportsAllAnnotations = step.annotations().any { it == "*" }
                return@associateWith if (
                    supportsAllAnnotations || elementsByAnnotation.isNotEmpty()
                ) {
                    step
                        .process(env, elementsByAnnotation, false)
                        .mapNotNull { (it.closestMemberContainer as? XTypeElement)?.qualifiedName }
                        .toSet()
                } else {
                    emptySet()
                }
            }
        // Store elements deferred by steps.
        elementsDeferredBySteps.clear()
        elementsDeferredBySteps.putAll(currentElementsDeferredByStep)
    }

    fun processLastRound(): List<String> {
        steps.forEach { step ->
            val stepDeferredElementsByAnnotation =
                getStepElementsByAnnotation(
                    step = step,
                    typeElementNames =
                        deferredElementNames +
                            elementsDeferredBySteps.getOrElse(step) { emptySet() }
                )
            val elementsByAnnotation =
                step
                    .annotations()
                    .mapNotNull { annotation ->
                        val annotatedElements =
                            stepDeferredElementsByAnnotation[annotation] ?: emptySet()
                        if (annotatedElements.isNotEmpty()) {
                            annotation to annotatedElements
                        } else {
                            null
                        }
                    }
                    .toMap()
            step.process(env, elementsByAnnotation, true)
        }
        // Return element names that were deferred until the last round, an error should be reported
        // for these, failing compilation. Sadly we currently don't have the mechanism to know if
        // the missing types were generated in the last round.
        return elementsDeferredBySteps.values.flatten()
    }

    /**
     * Get elements annotated with [step]'s annotations from the type element in [typeElementNames].
     *
     * Does not traverse type element members, so that if looking at `Outer` in the example below,
     * looking for `@X`, then `Outer`, `Outer.foo`, and `Outer.foo()` will be in the result, but
     * neither `Inner` nor its members will unless `Inner` is also in the [typeElementNames].
     *
     * ```
     * @X class Outer {
     *   @X Object foo;
     *   @X void foo() {}
     *   @X static class Inner {
     *     @X Object bar;
     *     @X void bar() {}
     *   }
     * }
     * ```
     */
    private fun getStepElementsByAnnotation(
        step: XProcessingStep,
        typeElementNames: Set<String>
    ): Map<String, Set<XElement>> {
        if (typeElementNames.isEmpty()) {
            return emptyMap()
        }
        val stepAnnotations = step.annotations()
        val elementsByAnnotation = mutableMapOf<String, MutableSet<XElement>>()
        fun putStepAnnotatedElements(element: XElement) =
            element
                .getAllAnnotations()
                .map { it.qualifiedName }
                .forEach { annotationName ->
                    if (stepAnnotations.contains(annotationName)) {
                        elementsByAnnotation
                            .getOrPut(annotationName) { mutableSetOf() }
                            .add(element)
                    }
                }
        typeElementNames
            .mapNotNull { env.findTypeElement(it) }
            .forEach { typeElement ->
                typeElement
                    .getEnclosedElements()
                    .filterNot { it.isTypeElement() }
                    .forEach { enclosedElement ->
                        if (enclosedElement is XExecutableElement) {
                            enclosedElement.parameters.forEach { putStepAnnotatedElements(it) }
                        }
                        putStepAnnotatedElements(enclosedElement)
                    }
                putStepAnnotatedElements(typeElement)
            }
        return elementsByAnnotation.withDefault { mutableSetOf() }
    }

    fun reportMissingElements(missingElementNames: List<String>) {
        missingElementNames.forEach { missingElementName ->
            env.messager.printMessage(
                kind = Diagnostic.Kind.ERROR,
                msg =
                    ("%s was unable to process '%s' because not all of its dependencies " +
                            "could be resolved. Check for compilation errors or a circular " +
                            "dependency with generated code.")
                        .format(processorClass.canonicalName, missingElementName),
            )
        }
    }
}
